package com.chukun.multi.masterslave.reusable;

import com.chukun.multi.utils.Debug;

import java.math.BigInteger;
import java.util.*;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;

public class ParallelPrimeGenerator {

    public static void main(String[] args) throws Exception {

        PrimeGeneratorService primeGeneratorService = new PrimeGeneratorService();

        Set<BigInteger> result = primeGeneratorService.generatePrime(10000000);

        System.out.println("generated "+result.size() +" prime ");
        System.out.println(result);

    }

}

class Range{
    public final int lowerBound;
    public final int upperBound;

    public Range(int lowerBound,int upperBound) {
        if (upperBound<lowerBound) {
            throw new IllegalArgumentException("upperBound should be greater than lowerBound!");
        }
        this.lowerBound = lowerBound;
        this.upperBound = upperBound;
    }

    @Override
    public String toString() {
        return "Range{" +
                "lowerBound=" + lowerBound +
                ", upperBound=" + upperBound +
                '}';
    }
}

class PrimeGeneratorService  extends AbstractMaster<Range, Set<BigInteger>,Set<BigInteger>> {

    public PrimeGeneratorService() {
        this.init();
    }

    @Override
    protected TaskDivideStrategy<Range> newTaskDivideStrategy(Object... params) {

        final int numberOfSlaves = slaves.size();
        final int originalTaskScale = Integer.parseInt(params[0].toString());
        final int subTaskScale = originalTaskScale/numberOfSlaves;
        final int subTasksCount = ((originalTaskScale%numberOfSlaves)==0)? numberOfSlaves:numberOfSlaves+1;

        /**
         * 定义任务的分解策略
         */
        TaskDivideStrategy<Range> tds = new TaskDivideStrategy<Range>() {

            int i = 0;
            @Override
            public Optional<Range> nextChunk() {
                int upperBound;
                if (i<subTasksCount) {
                    upperBound = i*subTaskScale;
                } else if (i == subTasksCount) {
                    upperBound = originalTaskScale;
                }else {
                    return Optional.empty();
                }
                int lowerBound = (i-1)*subTaskScale + 1;
                i++;
                return Optional.of(new Range(lowerBound,upperBound));
            }
        };
        return tds;
    }

    @Override
    protected Set<? extends SlaveSpec<Range, Set<BigInteger>>> createSlaves() {
        Set<PrimeGenerator> slaves = new HashSet<>();
        for (int i = 0; i < Runtime.getRuntime().availableProcessors(); i++) {
            slaves.add(new PrimeGenerator(new ArrayBlockingQueue<>(2)));
        }
        return slaves;
    }

    @Override
    protected Set<BigInteger> combineResult(Iterator<Future<Set<BigInteger>>> subResults) {

        Set<BigInteger> result = new TreeSet<>();
        while (subResults.hasNext()) {
            try {
                result.addAll(subResults.next().get());
            } catch (InterruptedException e) {
                e.printStackTrace();
            } catch (ExecutionException e) {
                Throwable cause = e.getCause();
                if (cause instanceof SubTaskFailureException) {
                    RetryInfo retryInfo = ((SubTaskFailureException)cause).getRetryInfo();
                    Object subTask = retryInfo.subTask;
                    Debug.info("failed subTask.. : "+subTask);
                }
            }
        }
        return result;
    }

    /**
     * 子类暴露一个具体服务的方法
     * @param upperBound
     * @return
     * @throws Exception
     */
    public Set<BigInteger> generatePrime(int upperBound) throws Exception {
        return this.service(upperBound);
    }


    /**
     * 质数生成器
     */
    private static class PrimeGenerator extends WorkerThreadSlave<Range,Set<BigInteger>> {

        public PrimeGenerator(BlockingQueue<Runnable> taskQueue) {
            super(taskQueue);
        }
        @Override
        protected Set<BigInteger> doProcess(Range range) throws Exception {
            Set<BigInteger> result = new TreeSet<>();
            BigInteger start = BigInteger.valueOf(range.lowerBound);
            BigInteger end = BigInteger.valueOf(range.upperBound);

            while ((start = start.nextProbablePrime()).compareTo(end) < 0) {
                result.add(start);
            }
            return result;
        }
    }
}
